library(slingshot)
library(scuttle)
library(scater)
library(scran)
library(ggraph)
library(tidySingleCellExperiment)

# simulate data representing a single lineage ------
set.seed(42)

means <- rbind(
  # non-DE genes
  matrix(rep(rep(c(0.1, 0.5, 1, 2, 3), each = 300), 100),
    ncol = 300, byrow = TRUE
  ),
  # early deactivation
  matrix(rep(exp(atan(((300:1) - 200) / 50)), 50), ncol = 300, byrow = TRUE),
  # late deactivation
  matrix(rep(exp(atan(((300:1) - 100) / 50)), 50), ncol = 300, byrow = TRUE),
  # early activation
  matrix(rep(exp(atan(((1:300) - 100) / 50)), 50), ncol = 300, byrow = TRUE),
  # late activation
  matrix(rep(exp(atan(((1:300) - 200) / 50)), 50), ncol = 300, byrow = TRUE),
  # transient
  matrix(rep(exp(atan(c((1:100) / 33, rep(3, 100), (100:1) / 33))), 50),
    ncol = 300, byrow = TRUE
  )
)
counts <- apply(means, 2, function(cell_means) {
  total <- rnbinom(1, mu = 7500, size = 4)
  rmultinom(1, total, cell_means)
})
rownames(counts) <- paste0("G", 1:750)
colnames(counts) <- paste0("c", 1:300)
sce <- SingleCellExperiment(assays = List(counts = counts))

# filter genes down to potential cell-type markers
# at least 3 reads in at least 10 cells
# apply(x, margin=1, fun) apply fun() to rows (gene) in x
geneFilter <- apply(assays(sce)$counts, 1, \(x)sum(x >= 3) >= 10)

sce <- sce[geneFilter, ]

# function of quantile normalization
FQnorm <- function(counts) {
  rk <- apply(counts, 2, rank, ties.method = "min")
  counts.sort <- apply(counts, 2, sort)
  refdist <- apply(counts.sort, 1, median)
  norm <- apply(rk, 2, \(r)refdist[r])
  rownames(norm) <- rownames(counts)
  norm
}
# add normalized expr as a new assay
assays(sce)$norm <- assays(sce)$counts |> FQnorm()

# scuttle::logNormCounts directly add logcounts assay
sce <- logNormCounts(sce)

logcounts(sce)[1:5,1:5]

## reduce dimension by PCA ------
pca <- sce |> assay('norm') |> log1p() |> t() |>
  prcomp(scale. = FALSE)

rd1 <- pca$x[, 1:2]

rd1 |>
  as_tibble() |>
  ggplot(aes(PC1, PC2)) +
  geom_point(alpha = .5)

# scater::runPCA or scran::fixedPCA
sce <- runPCA(sce)

sce |>
  ggplot(aes(PC1, PC2)) +
  geom_point(alpha = .5)

## reduce dimension by UMAP -----
rd2 <- sce |> assay('norm') |> log1p() |> t() |> umap()

colnames(rd2) <- c("UMAP1", "UMAP2")

rd2 |>
  as_tibble() |>
  ggplot(aes(UMAP1, UMAP2)) +
  geom_point(alpha = .5)

# add reduction results to SCE object
reducedDims(sce) <- SimpleList(PCA = rd1, UMAP = rd2)

# scater::runUMAP
sce <- runUMAP(sce)

sce |>
  ggplot(aes(UMAP1, UMAP2)) +
  geom_point(alpha = .5)

## scater::runTSNE ------
sce <- sce |> runTSNE()

sce |>
  ggplot(aes(TSNE1, TSNE2)) +
  geom_point(alpha = .5)

## clustering by Gaussian Mixture Modeling on pca -----------
cl1 <- Mclust(rd1)$classification
colData(sce)$GMM <- cl1

sce |>
  ggplot(aes(PC1, PC2, color = as.character(GMM))) +
  geom_point() +
  scale_color_brewer(type = 'qual', palette = 'Set1') +
  ggtitle('GMM clustering')

## clustering by k-means -------
cl2 <- kmeans(rd1, centers = 4)$cluster # choose k=4 arbitarily
colData(sce)$kmeans <- cl2

sce |>
  ggplot(aes(PC1, PC2, color = as.character(kmeans))) +
  geom_point() +
  scale_color_brewer(type = 'qual', palette = 'Set1') +
  ggtitle('Kmeans clustering on PCA')

sce |>
  ggplot(aes(UMAP1, UMAP2, color = as.character(kmeans))) +
  geom_point() +
  scale_color_brewer(type = 'qual', palette = 'Set1') +
  ggtitle('Kmeans clustering on UMAP')

sce |>
  ggplot(aes(TSNE1, TSNE2, color = as.character(kmeans))) +
  geom_point() +
  scale_color_brewer(type = 'qual', palette = 'Set1') +
  labs(title = 'Kmeans clustering on TSNE', color = 'kmeans')

## scran::clusterCells() by nearest-neighbor graph
nn.clusters <- clusterCells(sce, use.dimred = 'PCA')
table(nn.clusters)

sce$walktrap <- nn.clusters

## run slingshot with PCA & walktrap-----
sce_ss <- slingshot(sce, clusterLabels = "walktrap", reducedDim = "PCA")

# visualization: inferred lineage for the single-trajectory data with points colored by pseudotime
sce$slingPseudotime <- sce_ss$slingPseudotime_1

sce |>
  ggplot(aes(PC1, PC2)) +
  geom_point(aes(color = slingPseudotime)) +
  geom_path(data = slingCurves(sce_ss, as.df = TRUE)) +
  scale_color_viridis_c(option = 'turbo') +
  ggtitle('Pseudotime & curve on PCA')

## cluster-based minimum spanning tree ---------
ssgraph <- slingMST(sce_ss, as.df = TRUE) |>
  as_tibble()

sce |>
  ggplot(aes(PC1, PC2, color = as.character(walktrap))) +
  geom_point() +
  geom_path(data = ssgraph, color = 'black', linewidth = 2) +
  scale_color_brewer(type = 'qual', palette = 'Set1') +
  ggtitle('Kmeans clustering on PCA')

sce_ss |> slingMST() |>
  tidygraph::as_tbl_graph()

## Identify temporally dynamic genes -------
# BiocManager::install('tradeSeq')
library(tradeSeq)

# fit negative binomial GAM:
# Generalized Additive Models
sce_ss <- fitGAM(sce_ss)

# test for dynamic expression
ATres <- associationTest(sce_ss)

# pick most dynamic 250 genes
topgenes <- ATres |>
  slice_min(pvalue, n = 150, with_ties = FALSE) |>
  rownames()

pst.ord <- order(sce_ss$slingPseudotime_1, na.last = NA)

heatdata <- assays(sce)$logcounts[topgenes, pst.ord]
heatclus <- sce$kmeans[pst.ord]

heatmap(heatdata,
  Colv = NA,
  ColSideColors = brewer.pal(9, "Set1")[heatclus]
)

## project cells to existing trajectories -----
# simulate new cells in PCA space
newPCA <- reducedDim(sce, 'PCA') + rnorm(2*ncol(sce), sd = 2)

# project onto trajectory
newPTO <- sce_ss$slingshot |>
  slingshot::predict(newPCA)

newPTO |>
  as_tibble() |>
  select(1:2)

newPCA |>
  as_tibble() |>
  ggplot() +
  geom_point() +
  geom_point(data = sce, color = 'grey')

sce |>
  ggplot(aes(PC1, PC2)) +
  geom_point(color = 'grey') +
  geom_path(data = ssgraph, color = 'black', linewidth = 2) +
  geom_point(data = as_tibble(newPCA),
             aes(PC1, PC2, color = slingAvgPseudotime(newPTO))) +
  scale_color_viridis_c(option = 'turbo') +
  labs(color = 'Pseudotime of new\nprojected cells')

# bifurcation lineage example -----------
data("slingshotExample")
rd <- slingshotExample$rd
cl <- slingshotExample$cl

dim(rd) # 140 cells
table(cl) # in 5 clusters

## identify global lineage structure -----
lin1 <- getLineages(rd, cl, start.clus = "1")

bifr_mat <- bind_cols(rd, cl) |>
  setNames(c('dim.1','dim.2','cluster'))

lin1 |> slingMST(as.df = TRUE) |>
  ggplot(aes(dim.1, dim.2, group = Lineage)) +
  geom_path() +
  geom_point(size = 3) +
  geom_point(data = bifr_mat, aes(color = as.factor(cluster), group = NULL))

## force cluster 3 to be an endpoint -----
lin2 <- getLineages(rd, cl, start.clus = "1", end.clus = "3")

lin2 |> slingMST(as.df = TRUE) |>
  ggplot(aes(dim.1, dim.2, group = Lineage)) +
  geom_path() +
  geom_point(size = 3) +
  geom_point(data = bifr_mat, aes(color = as.factor(cluster), group = NULL))

## constructing smooth curves ------
lin1 |>
  getCurves() |>
  slingCurves(as.df = TRUE) |>
  arrange(Order) |>
  ggplot(aes(dim.1, dim.2, group = Lineage)) +
  geom_path() +
  geom_point(data = bifr_mat, aes(color = as.factor(cluster), group = NULL))

# For large datasets, we highgly recommend using the approx_points argument with slingshot (or getCurves)
# recommend a value of 100-200, hence the default value of 150
# restricting the number of unique points along the curve does not impose a similar limit on the number of unique pseudotime values, as demonstrated below
sce5 <- slingshot(sce,
  clusterLabels = "GMM", reducedDim = "PCA",
  approx_points = 5
)

# construct multiple trajectory by adjusting omega
sce5 <- slingshot(sce,
                  clusterLabels = "GMM", reducedDim = "PCA",
                  omega = TRUE)
